# deepseek_translator.py

import os
import re
import requests
import urllib3
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
from langdetect import detect, LangDetectException
from dotenv import load_dotenv  # 新增
import json
import time
import logging

# 禁用SSL警告
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

# 在类外加载 .env 文件
load_dotenv()

# 设置日志级别，避免过多输出
logger = logging.getLogger(__name__)

class BaseAITranslator:
    """AI翻译器基类"""
    
    def __init__(self, api_key=None, model="", source_lang="en", target_lang="zh-cn", endpoint=""):
        self.api_key = api_key
        self.model = model
        self.source_lang = source_lang
        self.target_lang = target_lang
        self.endpoint = endpoint
        
        # 创建带重试机制的session
        self.session = requests.Session()
        retry_strategy = Retry(
            total=3,
            backoff_factor=1,
            status_forcelist=[429, 500, 502, 503, 504],
        )
        adapter = HTTPAdapter(max_retries=retry_strategy)
        self.session.mount("http://", adapter)
        self.session.mount("https://", adapter)
        
        # 设置代理（如果配置了的话）
        proxies = {}
        if os.getenv("HTTP_PROXY"):
            proxies["http"] = os.getenv("HTTP_PROXY")
        if os.getenv("HTTPS_PROXY"):
            proxies["https"] = os.getenv("HTTPS_PROXY")
        if proxies:
            self.session.proxies.update(proxies)

    def detect_language(self, text):
        """检测文本语言"""
        try:
            return detect(text)
        except LangDetectException:
            return "unknown"

    def translate_text(self, text, max_retries=3):
        """
        使用 AI API 翻译单段文本

        :param text: 需要翻译的文本
        :param max_retries: 最大重试次数
        :return: 翻译结果字符串 或 原文
        """
        if not text.strip():
            return text

        # 自动检测语言，非源语言则跳过翻译
        try:
            detected_lang = self.detect_language(text)
            if detected_lang.lower() != self.source_lang.lower():
                return text
        except:
            pass  # 如果语言检测失败，继续尝试翻译

        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }

        system_prompt = f"You are a professional translator. Translate the following {self.source_lang} text to {self.target_lang}. Keep the meaning accurate and natural."

        payload = {
            "model": self.model,
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": text}
            ],
            "temperature": 0.5,
            "stream": False
        }

        for attempt in range(max_retries):
            try:
                response = self.session.post(
                    self.endpoint, 
                    headers=headers, 
                    json=payload, 
                    timeout=60,
                    verify=False  # 注意：生产环境中应设为True
                )
                response.raise_for_status()
                result = response.json()
                
                # 更安全地解析响应
                if "choices" in result and len(result["choices"]) > 0:
                    choice = result["choices"][0]
                    if "message" in choice and "content" in choice["message"]:
                        return choice["message"]["content"].strip()
                    else:
                        logger.debug(f"Unexpected API response structure: {list(result.keys())}")
                        return text
                else:
                    logger.debug(f"Unexpected API response structure: {list(result.keys())}")
                    return text
                    
            except requests.exceptions.RequestException as e:
                logger.warning(f"网络请求失败 (尝试 {attempt + 1}/{max_retries}): {str(e)}")
                if attempt == max_retries - 1:
                    return text
            except json.JSONDecodeError as e:
                logger.warning(f"JSON解析失败 (尝试 {attempt + 1}/{max_retries}): {str(e)}")
                logger.debug(f"响应内容: {response.text[:200]}")
                if attempt == max_retries - 1:
                    return text
            except KeyError as e:
                logger.warning(f"响应结构错误 (尝试 {attempt + 1}/{max_retries}): 缺少键 {str(e)}")
                logger.debug(f"响应内容: {response.text[:200]}")
                if attempt == max_retries - 1:
                    return text
            except Exception as e:
                logger.warning(f"翻译失败 (尝试 {attempt + 1}/{max_retries}): {str(e)}")
                if attempt == max_retries - 1:
                    return text

    def translate_subtitles(self, subtitles):
        """
        翻译字幕列表

        :param subtitles: 字幕数据列表，格式为：
                          [{'start': '00:00,160', 'end': '00:04,278', 'text': '...'}, ...]
        :return: 翻译后的字幕列表，新增 'translated' 字段
        """
        translated_subtitles = []

        for idx, sub in enumerate(subtitles):
            original_text = sub.get('text', '')
            translated_text = self.translate_text(original_text)

            translated_subtitles.append({
                'start': sub['start'],
                'end': sub['end'],
                'original': original_text,
                'translated': translated_text
            })

            # 每10个字幕打印一次进度，减少输出
            if (idx + 1) % 10 == 0 or idx == len(subtitles) - 1:
                logger.info(f"翻译进度: {idx + 1}/{len(subtitles)}")

            # 添加延迟以避免API速率限制
            time.sleep(0.1)

        return translated_subtitles


class DeepSeekTranslator(BaseAITranslator):
    def __init__(self, api_key=None, model="deepseek-chat", source_lang="en", target_lang="zh-cn"):
        """
        初始化 DeepSeek 翻译器

        :param api_key: DeepSeek API 密钥
        :param model: 使用的模型名称
        :param source_lang: 源语言代码（默认是英文 'en'）
        :param target_lang: 目标语言代码（默认是中文 'zh-cn'）
        """
        api_key = api_key or os.getenv("DEEPSEEK_API_KEY")
        if not api_key:
            raise ValueError("DeepSeek API key is required. Set it via environment variable DEEPSEEK_API_KEY")

        super().__init__(
            api_key=api_key,
            model=model,
            source_lang=source_lang,
            target_lang=target_lang,
            endpoint="https://api.deepseek.com/v1/chat/completions"
        )


class ChatGPTTranslator(BaseAITranslator):
    def __init__(self, api_key=None, model="gpt-3.5-turbo", source_lang="en", target_lang="zh-cn"):
        """
        初始化 ChatGPT 翻译器

        :param api_key: OpenAI API 密钥
        :param model: 使用的模型名称
        :param source_lang: 源语言代码（默认是英文 'en'）
        :param target_lang: 目标语言代码（默认是中文 'zh-cn'）
        """
        api_key = api_key or os.getenv("OPENAI_API_KEY")
        if not api_key:
            raise ValueError("OpenAI API key is required. Set it via environment variable OPENAI_API_KEY")

        super().__init__(
            api_key=api_key,
            model=model,
            source_lang=source_lang,
            target_lang=target_lang,
            endpoint="https://api.gptsapi.net/v1/chat/completions"
        )